import akshare as ak
import pandas as pd
import os

pd.set_option('display.max_columns', None)
pd.set_option('display.max_rows', None)
pd.set_option('display.width', 300)

GENERAL_DATA_PATH = './database'

class Data_Collection():
    def __init__(self):
        self.base_path = "./database"
        self.eft_list_df = pd.DataFrame()

    def get_ETF_whole_list(self, date=""):
        if len(date) == 8:
            fund_etf_spot_em_df = ak.fund_etf_spot_em(date=date)
        else:
            fund_etf_spot_em_df = ak.fund_etf_spot_em()
        # fund_etf_spot_em_df.to_csv(GENERAL_DATA_PATH + '/general_data/' + 'fund_etf_spot_simple.csv')
        fund_etf_spot_em_df = fund_etf_spot_em_df.sort_values(by="成交量", ascending=False)
        fund_etf_spot_em_df_temp = fund_etf_spot_em_df[fund_etf_spot_em_df["名称"].str.contains("沪深")]
        # print(fund_etf_spot_em_df_temp)
        self.eft_list_df = fund_etf_spot_em_df_temp
        return fund_etf_spot_em_df

    def get_single_ETF_info(self, code, start_date="1979-09-01 09:32:00", end_date="2222-01-01 09:32:00", period='5'):
        """
        fund_etf_hist_min_em 函数
        名称 	类型 	描述
        symbol 	str 	symbol='513500'; ETF 代码可以在 ak.fund_etf_spot_em() 中获取
        start_date 	str 	start_date="1979-09-01 09:32:00"; 日期时间; 默认返回所有数据
        end_date 	str 	end_date="2222-01-01 09:32:00"; 日期时间; 默认返回所有数据
        period 	str 	period='5'; choice of {'1', '5', '15', '30', '60'}; 其中 1 分钟数据返回近 5 个交易日数据且不复权
        adjust 	str 	adjust=''; choice of {'', 'qfq', 'hfq'}; '': 不复权, 'qfq': 前复权, 'hfq': 后复权, 其中 1 分钟数据返回近 5 个交易日数据且不复权
        """
        fund_etf_min_df = ak.fund_etf_hist_min_em(symbol=code, start_date=start_date, end_date=end_date, period=period)
        if not os.path.exists(GENERAL_DATA_PATH + '/etf_min_data/' + period):
            os.mkdir(GENERAL_DATA_PATH + '/etf_min_data/' + period)
        fund_etf_min_df.to_csv(GENERAL_DATA_PATH + '/etf_min_data/' + period + '/' + code + '_' + 'fund_etf_min_df.csv', index=False)
        return fund_etf_min_df

    def get_single_ETF_hist_info(self, code, start_date="20000101", end_date="20230104", period="daily"):
        """
        名称 	类型 	描述
        symbol 	str 	symbol='159707'; ETF 代码可以在 ak.fund_etf_spot_em() 中获取或查看东财主页
        period 	str 	period='daily'; choice of {'daily', 'weekly', 'monthly'}
        start_date 	str 	start_date='20000101'; 开始查询的日期
        end_date 	str 	end_date='20230104'; 结束查询的日期
        adjust 	str 	默认返回不复权的数据; qfq: 返回前复权后的数据; hfq: 返回后复权后的数据
        """
        fund_etf_hist_em_df = ak.fund_etf_hist_em(symbol=code, period="daily", start_date=start_date,
                                                  end_date=end_date, adjust="")
        if not os.path.exists(GENERAL_DATA_PATH + '/etf_hist_data/' + period):
            os.mkdir(GENERAL_DATA_PATH + '/etf_hist_data/' + period)
        fund_etf_hist_em_df.to_csv(GENERAL_DATA_PATH + '/etf_hist_data/' + period + '/' + code +'_fund_etf_hist_df.csv', index=False)


if __name__ == "__main__":
    collector = Data_Collection()

    # 获取ETF 参数列表
    etf_df = collector.get_ETF_whole_list()

    fund_etf_spot_em_df_temp = etf_df[etf_df["名称"].str.contains("沪深")]

    code_list = fund_etf_spot_em_df_temp["代码"]

    # 获取个体 ETF 日数据
    # collector.get_eft_info(code)
    for code in code_list:
        # 获取实时数据并保存
        collector.get_single_ETF_info(code, period="1")
        # 获取历史数据并保存
        # collector.get_single_ETF_hist_info(code)

